#include <cstdio>
#include <memory>
#include <string>

#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/image_frame_opencv.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/opencv_highgui_inc.h"
#include "mediapipe/framework/port/opencv_imgproc_inc.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/gpu/gl_calculator_helper.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/gpu/gpu_shared_data_internal.h"

// 示例配置：将图像上传到 GPU 并调整大小
constexpr char kConfig[] = R"(
input_stream: "input"
output_stream: "output"

node {
  calculator: "GpuBufferToImageFrameCalculator"
  input_stream: "input"
  output_stream: "output_cpu"
}

node {
  calculator: "ImageFrameToGpuBufferCalculator"
  input_stream: "output_cpu"
  output_stream: "output"
}
)";

namespace mediapipe {

absl::Status RunMediaPipeGPU() {
  // 初始化 MediaPipe
  MP_RETURN_IF_ERROR(mediapipe::InitializeMediaPipe());

  // 创建 GPU 资源
  ASSIGN_OR_RETURN(auto gpu_resources, mediapipe::GpuResources::Create());
  GlCalculatorHelper gpu_helper;
  gpu_helper.InitializeForTest(gpu_resources.get());

  // 解析配置
  CalculatorGraphConfig config =
      ParseTextProtoOrDie<CalculatorGraphConfig>(kConfig);

  // 创建计算图
  CalculatorGraph graph;
  MP_RETURN_IF_ERROR(graph.Initialize(config));
  MP_RETURN_IF_ERROR(gpu_helper.Setup(&graph));

  // 获取输入输出流
  ASSIGN_OR_RETURN(OutputStreamPoller poller,
                   graph.AddOutputStreamPoller("output"));
  ASSIGN_OR_RETURN(OutputStreamPoller input_poller,
                   graph.AddOutputStreamPoller("output_cpu"));

  // 启动计算图
  MP_RETURN_IF_ERROR(graph.StartRun({}));

  // 创建测试图像（512x512 红色图像）
  const int width = 512;
  const int height = 512;
  auto input_frame = std::make_shared<ImageFrame>(
      ImageFormat::SRGB, width, height, ImageFrame::kGlDefaultAlignmentBoundary);
  cv::Mat input_mat = formats::MatView(input_frame.get());
  input_mat.setTo(cv::Scalar(0, 0, 255)); // BGR 格式的红色

  // 发送数据到计算图
  MP_RETURN_IF_ERROR(gpu_helper.RunInGlContext([&]() -> absl::Status {
    // 将图像转换为 GPU Buffer
    GpuBuffer gpu_buffer = gpu_helper.CreateSourceTexture(*input_frame).GetFrame<GpuBuffer>();
    
    // 发送到输入流
    MP_RETURN_IF_ERROR(graph.AddPacketToInputStream(
        "input", Adopt(gpu_buffer.release()).At(Timestamp(0))));
    
    return absl::OkStatus();
  }));

  // 获取处理结果
  mediapipe::Packet packet;
  if (poller.Next(&packet)) {
    auto& output_gpu = packet.Get<GpuBuffer>();
    
    // 将 GPU 数据转换回 CPU
    std::shared_ptr<const ImageFrame> output_frame;
    MP_RETURN_IF_ERROR(gpu_helper.RunInGlContext([&]() -> absl::Status {
      output_frame = output_gpu.GetReadView<ImageFrame>(gpu_helper.gl_context());
      return absl::OkStatus();
    }));

    // 显示结果
    cv::Mat output_mat = formats::MatView(output_frame.get());
    cv::imshow("Output", output_mat);
    cv::waitKey(0);
  }

  // 关闭计算图
  return graph.Close();
}

}  // namespace mediapipe

int main() {
  if (mediapipe::RunMediaPipeGPU().ok()) {
    std::cout << "GPU 功能验证成功!" << std::endl;
    return 0;
  }
  std::cerr << "GPU 功能验证失败!" << std::endl;
  return 1;
}